import chromadb
from chromadb.utils import embedding_functions 
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
from sentence_transformers import SentenceTransformer
from common.config import config

#model = SentenceTransformer(model_name_or_path=model_path)
model = SentenceTransformer(model_name_or_path='C:/Users/Administrator/.cache/modelscope/hub/models/Jerry0/text2vec-base-chinese')
#embedding_functions.HuggingFaceEmbeddingFunction(model_id='C:/Users/Administrator/.cache/modelscope/hub/models/Jerry0/text2vec-base-chinese')

sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name='C:/Users/Administrator/.cache/modelscope/hub/models/Jerry0/text2vec-base-chinese')


class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        embeddings = [list(model.encode(text,normalize_embeddings=True).astype(float)) for text in texts]
        return embeddings
ebfun=MyEmbeddingFunction()

class adbsvc:
    
    def __init__(self):
        self.chroma_client = chromadb.PersistentClient(path=config["adb"]["db_path"])
    def getcli(self):
        return self.chroma_client 
def cli():
    #chroma_client = chromadb.PersistentClient(path="D:/ai/adbdisk")
    chroma_client = chromadb.PersistentClient(path=config["adb"]["db_path"])
    return chroma_client
#,embedding_function=MyEmbeddingFunction()
def getCollOk0(chroma_client,coll):
    collection = chroma_client.get_or_create_collection(coll,  embedding_function=ebfun)
    return collection

def getColl(chroma_client,coll):
    collection = chroma_client.get_or_create_collection(coll,  embedding_function=sentence_transformer_ef)
    return collection   

def cli0():
    chroma_client = chromadb.PersistentClient(path="D:/ai/adbdisk")
    return chroma_client
 
 
def add(coll,_documents,metadatas,idd):
    chroma_client=cli();
    #collection = chroma_client.get_or_create_collection(coll)
    collection=getColl(chroma_client,coll)
    collection.add(documents=_documents,metadatas=metadatas,ids=idd)

def qry(coll,msg,n_results:2):
    chroma_client=cli();
    #collection = chroma_client.get_or_create_collection(coll)
    collection=getColl(chroma_client,coll)
    results=collection.query(
    query_texts=msg,
    n_results=n_results)
    return results
 
 
